import tensorflow as tf
import time
import numpy as np

def test_mps_support():
    """
    测试TensorFlow对MPS的支持
    """
    print("TensorFlow版本:", tf.__version__)
    print("可用的GPU设备:", tf.config.list_physical_devices('GPU'))
    
    # 检查MPS支持
    if tf.config.list_physical_devices('GPU'):
        print("MPS支持可用")
        return True
    else:
        print("MPS支持不可用，将使用CPU")
        return False

def create_model():
    """
    创建一个简单的神经网络模型
    """
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(512, activation='relu', input_shape=(784,)),
        tf.keras.layers.Dense(256, activation='relu'),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
    return model

def create_data():
    """
    创建测试数据
    """
    # 生成模拟MNIST数据集大小的数据
    x_train = np.random.random((60000, 784)).astype(np.float32)
    y_train = np.random.randint(0, 10, (60000,))
    
    x_test = np.random.random((10000, 784)).astype(np.float32)
    y_test = np.random.randint(0, 10, (10000,))
    
    return (x_train, y_train), (x_test, y_test)

def benchmark_device(device_name, device):
    """
    在指定设备上运行基准测试
    """
    print(f"\n在 {device_name} 上运行基准测试...")
    
    # 创建模型和数据
    with tf.device(device):
        model = create_model()
        (x_train, y_train), (x_test, y_test) = create_data()
        
        # 训练模型并计时
        start_time = time.time()
        model.fit(x_train, y_train, epochs=5, batch_size=128, verbose=1)
        training_time = time.time() - start_time
        
        # 评估模型
        start_time = time.time()
        test_loss, test_acc = model.evaluate(x_test, y_test, verbose=0)
        evaluation_time = time.time() - start_time
        
    print(f"{device_name} 训练时间: {training_time:.2f} 秒")
    print(f"{device_name} 评估时间: {evaluation_time:.2f} 秒")
    print(f"{device_name} 测试准确率: {test_acc:.4f}")
    
    return training_time, evaluation_time, test_acc

def main():
    """
    主函数
    """
    print("TensorFlow MPS 性能测试")
    print("=" * 30)
    
    # 检查MPS支持
    mps_available = test_mps_support()
    
    # 创建基准测试结果存储
    results = {}
    
    # 在CPU上运行基准测试
    cpu_training_time, cpu_evaluation_time, cpu_accuracy = benchmark_device("CPU", "/CPU:0")
    results["CPU"] = {
        "training_time": cpu_training_time,
        "evaluation_time": cpu_evaluation_time,
        "accuracy": cpu_accuracy
    }
    
    # 如果MPS可用，在GPU上运行基准测试
    if mps_available:
        gpu_training_time, gpu_evaluation_time, gpu_accuracy = benchmark_device("GPU(MPS)", "/GPU:0")
        results["GPU"] = {
            "training_time": gpu_training_time,
            "evaluation_time": gpu_evaluation_time,
            "accuracy": gpu_accuracy
        }
        
        # 计算加速比
        training_speedup = cpu_training_time / gpu_training_time
        evaluation_speedup = cpu_evaluation_time / gpu_evaluation_time
        
        print("\n性能对比:")
        print(f"训练加速比: {training_speedup:.2f}x")
        print(f"评估加速比: {evaluation_speedup:.2f}x")
    else:
        print("\n跳过GPU测试，因为MPS不可用")

if __name__ == "__main__":
    main()